#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <numeric>

using namespace std;


class Solution {
public:
    using int_pair = pair<int, int>;
    static bool cmp(int_pair a, int_pair b) {
        return a.first < b.first;
    }

    vector<int> sort_indices(const vector<int>& nums) {
        vector<int> res(nums.size());
        iota(res.begin(), res.end(), 0);
        // for (size_t i = 0; i != idx.size(); ++i) idx[i] = i;  相当于这个

        stable_sort(res.begin(), res.end(), 
            [&nums](int i1, int i2) { return nums[i1] < nums[i2]; });
        
        return res;
    }

    int findMaximizedCapital(int k, int w, vector<int>& profits, vector<int>& capital) {
        int n = profits.size();

        // Sort by capital
        vector<int> is = sort_indices(capital);
        stable_sort(capital.begin(), capital.end());  // capital sorted
        priority_queue<int_pair, vector<int_pair>, decltype(&cmp)> pq(cmp);  // MaxHeap

        //* All capitals in the heap < w
        int i = 0;
        while (i < n && capital[i] <= w) {
            pq.emplace(profits[is[i]], capital[i]);
            i++;
        }

        int cnt = 0;
        while (cnt < k) {
            if (!pq.empty()) {
                w += get<0>(pq.top()); pq.pop();
                cnt++;
                while (i < n && capital[i] <= w) {
                    pq.emplace(profits[is[i]], capital[i]);
                    i++;
                }

            } else if (i >= n || capital[i] > w) break;
            //* 已经包括全部情况了
        }

        return w;
    }
};


int main() {
    Solution so;
    vector<int> profits{2,4,5,8,0,1,6,3,9,12,6,14,2,3,3};
    vector<int> capital{1,4,5,6,8,2,4,5,5,7,0,12,1,1,6};
    vector<int> profits1{40,2};
    vector<int> capital1{18,3};
    cout << so.findMaximizedCapital(6, 6, profits1, capital1) << endl;

    return 0;
}
